# encoding=utf-8
# https://docs.openvino.ai/2024/openvino-workflow/model-preparation/convert-model-pytorch.html

import torch
import torch.nn as nn
from torchvision import models
import openvino as ov

model = models.densenet121(pretrained=True)
model.classifier = nn.Linear(model.classifier.in_features, 5)
model.load_state_dict(torch.load("state.pth"))
model.eval()

# convert from pytorch
print("ov version", ov.get_version())
ov_model = ov.convert_model(model,
                            example_input=torch.rand(1, 3, 224, 224),
                            input=("x", [1, 3, 224, 224]),
                            verbose=True)
ov.save_model(ov_model, "export_dense121_cpu.xml", compress_to_fp16=False)
# ov.serialize(ov_model, bin_path="export_dense121_cpu.bin", xml_path="export_dense121_cpu.xml")